perf(index): use dot distance for cosine in FlatFloatStorage - #7777
perf(index): use dot distance for cosine in FlatFloatStorage#7777clearlove10-c wants to merge 1 commit into
Conversation
For IVF indexes with cosine distance, vectors are already normalized by the IVF transform pipeline (NormalizeTransformer), so cosine distance is mathematically equivalent to dot distance over normalized vectors. However, cosine still computes the L2 norm of each vector during distance calculation — redundant work that dot avoids. This converts Cosine to Dot in FlatFloatStorage::try_from_batch, consistent with how PQ and RQ already convert Cosine to L2 for the same reason. Adds tests verifying: - Cosine is converted to Dot - L2 and Dot are preserved unchanged - Partition serde round-trip adapts to the conversion
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughFlat vector storage now normalizes cosine distance to dot distance for IVF-normalized vectors. Tests cover conversion, preservation of other distance types, cosine/dot equivalence, and serialized codec restoration. ChangesDistance normalization
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related issues
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
| // transform pipeline, so cosine distance is equivalent to dot distance | ||
| // over normalized vectors. Use dot to avoid redundant norm computation. | ||
| let distance_type = match distance_type { | ||
| DistanceType::Cosine => DistanceType::Dot, |
There was a problem hiding this comment.
Converting Cosine to Dot here changes observable Float16 distance and range behavior even when the vectors pass through the real IVF normalization pipeline. Float16 rounding leaves the normalized norm different from 1: on this head [7, 47, 13] normalizes to [0.14221191, 0.9550781, 0.26416016]; the existing cosine kernel returns +1.19e-7 for the self-distance, while this storage returns -0.002179027. FlatIndex drops distances below lower_bound, so a cosine range query with lower_bound=0 now excludes the vector's own self-match. The same rounding can also reverse near-neighbor order.
Could we keep the logical metric as Cosine and cache the query norm in FlatDistanceCal (and, if benchmarks justify it, derived per-vector norms in storage)? That removes redundant norm scans without requiring exact unit norms.
Reproducer run on 7ee1282
#[test]
fn normalized_f16_cosine_keeps_self_match_at_zero_lower_bound() {
let values = Float16Array::from(vec![
f16::from_f32(7.0),
f16::from_f32(47.0),
f16::from_f32(13.0),
]);
let raw = FixedSizeListArray::try_new_from_values(values, 3).unwrap();
let normalized = lance_linalg::kernels::normalize_fsl(&raw).unwrap();
let query = normalized.value(0);
let batch = make_flat_test_batch(normalized);
let storage = FlatFloatStorage::try_from_batch(
batch,
&FlatMetadata { dim: 3 },
DistanceType::Cosine,
None,
)
.unwrap();
let distance = storage.dist_calculator(query, 0.0).distance(0);
assert!(
distance >= 0.0,
"lower_bound=0 would drop self-match: {distance}"
);
}Current result: lower_bound=0 would drop self-match: -0.0021790266.
2 participants
Problem
For IVF indexes with cosine distance, vectors are already normalized
by the IVF transform pipeline (NormalizeTransformer), so cosine
distance is mathematically equivalent to dot distance over normalized
vectors. However, cosine still computes the L2 norm of each vector
during distance calculation — redundant work that dot avoids.
Solution
Convert Cosine to Dot in
FlatFloatStorage::try_from_batch, consistentwith how PQ and RQ already convert Cosine to L2 for the same reason.
Benchmark
scripts below
my_ann_bench.py
Tests
test_try_from_batch_converts_cosine_to_dot— verifies Cosine → Dottest_try_from_batch_keeps_non_cosine_distance_types— verifies L2/Dot unchangedSummary by CodeRabbit